import pickle
import numpy as np
import scipy.io.wavfile as wf
from flask import Flask, request, jsonify
from flask_cors import CORS
import os
from scipy.fftpack import fft

class AudioRecognizer:
    def __init__(self, model_path):
        self.names = ['abnormal_treble（不正常高音）', 'normal_bass（正常低音）', 'normal_bass_inside_the_carriage（正常车厢内低音）', 'normal_treble（正常高音）', 'normal_treble_inside_the_carriage（正常车厢内高音）', 'resonant_bass_inside_the_carriage（共振车厢内低音）', 'resonant_treble_inside_the_carriage（共振车厢内高音）']
        with open(model_path, 'rb') as file:
            self.model = pickle.load(file)

    def analyze_frequency(self, file):
        sample_rate, data = wf.read(file)
        if data.ndim > 1:
            data = data[:, 0]
        fft_out = fft(data)
        freqs = np.fft.fftfreq(len(fft_out))
        idx = np.argmax(np.abs(fft_out))
        freq = freqs[idx]
        freq_in_hertz = abs(freq * sample_rate)
        return freq_in_hertz

    def predict(self, file_path):
        frequency = self.analyze_frequency(file_path)
        features = np.array([[frequency]])
        prediction = self.model.predict(features)
        return self.names[prediction[0]]

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'uploads'
os.makedirs('uploads', exist_ok=True)

CORS(app)

def allowed_file(filename):
    ALLOWED_EXTENSIONS = set(['aac', 'wav', 'mp3'])
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

@app.route('/upload', methods=['POST'])
def upload_audio():
    if 'file' not in request.files:
        return jsonify(status='error', message='No file part'), 400
    file = request.files['file']

    if file.filename == '':
        return jsonify(status='error', message='No selected file'), 400
    if file and allowed_file(file.filename):
        filename = os.path.join(app.config['UPLOAD_FOLDER'], file.filename)
        file.save(filename)
        recognizer = AudioRecognizer(model_path="model.pickle")
        # recognizer = AudioRecognizer(model_path="D:/graduation design/model_deployment/voc_svm/voc_svm/model_mfc.pickle")
        predicted_label = recognizer.predict(filename)
        print("正在识别,请稍后......")
        print(f"识别结果: {predicted_label}")  # 来打印结果
        return jsonify(status='success', result=predicted_label), 200
    else:
        return jsonify(status='error', message='File type not allowed'), 400

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=True)